// API Hook
// Thanks to Haifeng[RCT]
// 

#include "stdafx.h"
#include "HookApi.h"
#include "WinsockHook.h"

#define MAXHOOKCOUNT 100

DWORD GetBeginAddr();
DWORD GetEndAddr();

#pragma data_seg("HookData")
INT g_lstThreadID[MAXHOOKCOUNT] = {0};
HHOOK g_lstHhook[MAXHOOKCOUNT] = {0};
#pragma data_seg()

extern HMODULE m_hModule;

HHOOK m_currentHhook = NULL;
DWORD m_currentProcessID;
CWinsockHook *m_pWinsockHook = NULL;

#pragma comment(linker, "/SECTION:HookStub,RWE")

#define ALLOCATE_HookStub ALLOCATE(HookStub)

#pragma code_seg("HookStub")
#pragma optimize("gsy",on)
ALLOCATE_HookStub HOOKENVIRONMENT pEnv={0};
NAKED DWORD GetBeginAddr()
{
	__asm
	{
		call next
next:
		pop eax
			sub eax,5
			ret
	}
}
NAKED DWORD GetDelta()
{
	__asm
	{
		call next
next:
		pop eax
			sub eax,offset next
			ret
	}
}
NAKED void NewStub()
{
	__asm
	{	
		jmp next
back:
		_emit 0xE9
			NOP
			NOP
			NOP
			NOP
next:
		push [esp]
		push [esp]
		push eax
			call GetDelta
			call GetBeginAddr
			lea eax,[eax-0x20]	//size of HOOKENVIRONMENT
		mov dword ptr [esp+0xC],eax
			pop eax
			jmp back
	}
}
NAKED DWORD GetEndAddr()
{
	__asm
	{
		call next
next:
		pop eax
			push eax
			sub DWORD ptr [esp],5
			pop eax
			ret
	}
}
#pragma optimize("",off)
#pragma code_seg()


DWORD __stdcall GetOpCodeSize(BYTE* iptr0)
{
	BYTE* iptr = iptr0;

	DWORD f = 0;

prefix:
	BYTE b = *iptr++;

	f |= table_1[b];

	if (f&C_FUCKINGTEST)
		if (((*iptr)&0x38)==0x00)   // ttt
			f=C_MODRM+C_DATAW0;       // TEST
		else
			f=C_MODRM;                // NOT,NEG,MUL,IMUL,DIV,IDIV

	if (f&C_TABLE_0F)
	{
		b = *iptr++;
		f = table_0F[b];
	}

	if (f==C_ERROR)
	{
		//printf("error in %02X\n",b);
		return C_ERROR;
	}

	if (f&C_PREFIX)
	{
		f&=~C_PREFIX;
		goto prefix;
	}

	if (f&C_DATAW0) if (b&0x01) f|=C_DATA66; else f|=C_DATA1;

	if (f&C_MODRM)
	{
		b = *iptr++;
		BYTE mod = b & 0xC0;
		BYTE rm  = b & 0x07;
		if (mod!=0xC0)
		{
			if (f&C_67)         // modrm16
			{
				if ((mod==0x00)&&(rm==0x06)) f|=C_MEM2;
				if (mod==0x40) f|=C_MEM1;
				if (mod==0x80) f|=C_MEM2;
			}
			else                // modrm32
			{
				if (mod==0x40) f|=C_MEM1;
				if (mod==0x80) f|=C_MEM4;
				if (rm==0x04) rm = (*iptr++) & 0x07;    // rm<-sib.base
				if ((rm==0x05)&&(mod==0x00)) f|=C_MEM4;
			}
		}
	} // C_MODRM

	if (f&C_MEM67)  if (f&C_67) f|=C_MEM2;  else f|=C_MEM4;
	if (f&C_DATA66) if (f&C_66) f|=C_DATA2; else f|=C_DATA4;

	if (f&C_MEM1)  iptr++;
	if (f&C_MEM2)  iptr+=2;
	if (f&C_MEM4)  iptr+=4;

	if (f&C_DATA1) iptr++;
	if (f&C_DATA2) iptr+=2;
	if (f&C_DATA4) iptr+=4;

	return iptr - iptr0;
}

PHOOKENVIRONMENT __stdcall InstallHook(PCHAR DllName,PCHAR ApiName,PVOID HookProc)
{
	HMODULE DllHandle;
	PVOID ApiEntry;
	int ReplaceCodeSize;
	DWORD oldpro;
	DWORD SizeOfStub;
	DWORD delta;
	DWORD RetSize =0;

	PHOOKENVIRONMENT pHookEnv;

	if (HookProc == NULL)
	{
		return NULL;
	}

	DllHandle = GetModuleHandleA(DllName);
	if (DllHandle == NULL)
		DllHandle = LoadLibraryA(DllName);
	if (DllHandle == NULL)
		return NULL;

	ApiEntry = GetProcAddress(DllHandle,ApiName);
	if (ApiEntry == NULL) return NULL;

	ReplaceCodeSize = GetOpCodeSize((BYTE*)ApiEntry);

	while (ReplaceCodeSize < 5)
		ReplaceCodeSize += GetOpCodeSize((BYTE*)((DWORD)ApiEntry + (DWORD)ReplaceCodeSize));

	if (ReplaceCodeSize > 16) return NULL;

	SizeOfStub = sizeof(HOOKENVIRONMENT) + GetEndAddr() - GetBeginAddr();

	pHookEnv = (PHOOKENVIRONMENT)VirtualAlloc(NULL,SizeOfStub,MEM_COMMIT,PAGE_EXECUTE_READWRITE);
	memset((void*)&pEnv,0x90,sizeof(pEnv));
	CopyMemory(pHookEnv,(PVOID)&pEnv,sizeof(HOOKENVIRONMENT));
	CopyMemory((void*)((DWORD)pHookEnv+sizeof(HOOKENVIRONMENT)),(PVOID)GetBeginAddr(),GetEndAddr()-GetBeginAddr());
	CopyMemory((void*)pHookEnv,(void*)&pEnv,sizeof(pEnv.savebytes));
	CopyMemory(pHookEnv->savebytes,ApiEntry,ReplaceCodeSize);

	pHookEnv->OrgApiAddr = ApiEntry;
	pHookEnv->SizeOfReplaceCode = ReplaceCodeSize;

	pHookEnv->jmptoapi[0]=0xE9;
	*(DWORD*)(&pHookEnv->jmptoapi[1]) = (DWORD)ApiEntry + ReplaceCodeSize - ((DWORD)pHookEnv->jmptoapi + 5);

	if (!VirtualProtect(ApiEntry,ReplaceCodeSize,PAGE_EXECUTE_READWRITE,&oldpro))
		return FALSE;

	delta = (DWORD)pHookEnv - ((DWORD)GetBeginAddr()-sizeof(HOOKENVIRONMENT));

	*(DWORD*)(&JMPGate[1]) = ((DWORD)NewStub + delta) - ((DWORD)ApiEntry + 5);

	WriteProcessMemory(GetCurrentProcess(), ApiEntry, JMPGate, sizeof(JMPGate),&RetSize);

	if (!VirtualProtect(ApiEntry,ReplaceCodeSize,oldpro,&oldpro))
		return FALSE;

	*(DWORD*)((DWORD)NewStub + delta + 3) = (DWORD)HookProc - ((DWORD)NewStub + delta + 3 + 4);

	return pHookEnv;
}

BOOL __stdcall UninstallHook(PHOOKENVIRONMENT pHookEnv)
{
	DWORD oldpro;
	DWORD RetSize;

	if(IsBadReadPtr((const void*)pHookEnv,sizeof(HOOKENVIRONMENT)))
		return FALSE;

	if(!VirtualProtect(pHookEnv->OrgApiAddr,pHookEnv->SizeOfReplaceCode,PAGE_EXECUTE_READWRITE,&oldpro))
		return FALSE;
	WriteProcessMemory(GetCurrentProcess(),pHookEnv->OrgApiAddr,pHookEnv->savebytes,pHookEnv->SizeOfReplaceCode,&RetSize);

	if(!VirtualProtect(pHookEnv->OrgApiAddr,pHookEnv->SizeOfReplaceCode,oldpro,&oldpro))
		return FALSE;

	VirtualFree((LPVOID)pHookEnv,0,MEM_RELEASE);
	return TRUE;
}

VOID AddHook(DWORD dwThreadID, HHOOK hHook)
{
	// TODO: Synchronization
	for(int i = 0; i < MAXHOOKCOUNT; ++i)
	{
		if(g_lstThreadID[i] == NULL && g_lstHhook[i] == NULL)
		{
			g_lstThreadID[i] = dwThreadID;
			g_lstHhook[i] = hHook;
			break;
		}
	}
}

HHOOK GetHook(DWORD dwProcessID)
{
	for(int i = 0; i < MAXHOOKCOUNT; ++i)
	{
		if(g_lstThreadID[i] == dwProcessID)
		{
			return g_lstHhook[i];
		}
	}
	return NULL;
}

VOID RemoveHook(DWORD dwProcessID)
{
	// TODO: Synchronization
	for(int i = 0; i < MAXHOOKCOUNT; ++i)
	{
		if(g_lstThreadID[i] == dwProcessID)
		{
			g_lstThreadID[i] = NULL;
			g_lstHhook[i] = NULL;
			break;
		}
	}
}

static HMODULE ModuleFromAddress(PVOID pv) 
{
	MEMORY_BASIC_INFORMATION mbi;
	if(::VirtualQuery(pv, &mbi, sizeof(mbi)) != 0)
	{
		return (HMODULE)mbi.AllocationBase;
	}
	else
	{
		return NULL;
	}
}

static LRESULT WINAPI GetMsgProc(int code, WPARAM wParam, LPARAM lParam) 
{
	return ::CallNextHookEx(m_currentHhook, code, wParam, lParam);
}

BOOL WINAPI SetHook(BOOL bInstall, DWORD dwProcessID, DWORD /* dwThreadId */)
{
	BOOL bOk = 0;
	HHOOK hHook = NULL;
	if(bInstall)
	{
		// hHook = SetWindowsHookEx(WH_GETMESSAGE, GetMsgProc, ModuleFromAddress(GetMsgProc), dwThreadId);
		// Create a remote thread and hook it
		TCHAR strModulePath[2000] = {0};
		GetModuleFileName(m_hModule, strModulePath, 2000);
		HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, false, dwProcessID);
		FARPROC farLoadLibrary = GetProcAddress(GetModuleHandle(_T("Kernel32.dll")), "LoadLibraryW");
		LPVOID lpDllAddr = VirtualAllocEx(hProcess, NULL, wcslen(strModulePath) * sizeof(TCHAR), MEM_COMMIT, PAGE_READWRITE); 
		if(lpDllAddr != NULL)
		{
			if(WriteProcessMemory(hProcess, lpDllAddr, strModulePath, wcslen(strModulePath) * sizeof(TCHAR), NULL))
			{
				HANDLE hT = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)farLoadLibrary, lpDllAddr, 0, NULL);   
				CloseHandle(hT);
				CloseHandle(hProcess);
				bOk = 1;
			}
		}
		AddHook(dwProcessID, (HHOOK)1);
// 		if(hHook != NULL) 
// 		{
// 			bOk = 1;
// 			AddHook(dwThreadId, hHook);
// 		}
// 		else
// 		{
// 			bOk=0;
// 		}
	} 
	else 
	{
// 		UnInstallHookApi(pHookEnv_send);
// 		UnInstallHookApi(pHookEnv_sendto);
// 		UnInstallHookApi(pHookEnv_recv);
// 		UnInstallHookApi(pHookEnv_recvfrom);
// 		UnInstallHookApi(pHookEnv_WSARecv);
// 		UnInstallHookApi(pHookEnv_WSARecvFrom);
		delete m_pWinsockHook;
		bOk = UnhookWindowsHookEx(GetHook(dwProcessID));
		RemoveHook(dwProcessID);
	}
	return bOk;
}